#include <Windows.h>
#include <tchar.h>
#include "resource.h"
#include <Tlhelp32.h>
#include <commctrl.h>
#include <vector>
using std::vector;

UINT RetrieveProcess(vector<PROCESSENTRY32>& vtProcessList)
{
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (INVALID_HANDLE_VALUE != hSnapshot)
	{
		PROCESSENTRY32 pe32 = {0};
		pe32.dwSize	= sizeof(pe32);
		if (Process32First(hSnapshot, &pe32))
		{
			do {
				vtProcessList.push_back(pe32);
			} while (Process32Next(hSnapshot, &pe32));
		}
	}
	CloseHandle(hSnapshot);

	return vtProcessList.size();
}

UINT RetrieveModules(vector<MODULEENTRY32>& vtModuleList, HANDLE hProcess)
{
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, (DWORD)hProcess);
	if (INVALID_HANDLE_VALUE != hSnapshot)
	{
		MODULEENTRY32 me = {0};
		me.dwSize	= sizeof(me);
		if (Module32First(hSnapshot, &me))
		{
			do {
				vtModuleList.push_back(me);
			} while (Module32Next(hSnapshot, &me));
		}
	}
	CloseHandle(hSnapshot);

	return 0;
}

int CALLBACK DlgProc(HWND hDlg, UINT nMsg, WPARAM wParam, LPARAM lParam)
{
	static HWND hListView = NULL;
	switch (nMsg)
	{
		case WM_INITDIALOG:
		{
			SetWindowText(hDlg, _T("Demo0075"));

			UINT nWidth		= GetSystemMetrics(SM_CXSCREEN);
			UINT nHeight	= GetSystemMetrics(SM_CYSCREEN);
			RECT rtClient;
			GetClientRect(hDlg, &rtClient);
			SetWindowPos(hDlg, NULL, (nWidth - (rtClient.right - rtClient.left)) / 2, (nHeight - (rtClient.bottom - rtClient.top)) / 2, 0, 0, SWP_NOSIZE);
			
			hListView = GetDlgItem(hDlg, IDC_LVPROCESS);
			
			LVCOLUMN lvCol = {0};
			lvCol.mask		= LVCF_WIDTH|LVCF_TEXT;
			lvCol.pszText	= _T("File Name");
			lvCol.cx		= 250;
			ListView_InsertColumn(hListView, 0, &lvCol);
			lvCol.pszText	= _T("Thread Count");
			lvCol.cx		= 80;
			ListView_InsertColumn(hListView, 1, &lvCol);
			lvCol.pszText	= _T("Process ID");
			lvCol.cx		= 80;
			ListView_InsertColumn(hListView, 2, &lvCol);
			ListView_SetExtendedListViewStyle(hListView, LVS_EX_GRIDLINES|LVS_EX_FULLROWSELECT);

			SendMessage(hDlg, WM_COMMAND, MAKEWPARAM(IDC_BTNREFRESH, BM_CLICK), 0);
			break;
		}
		case WM_NOTIFY:
		{
			switch (((LPNMHDR)lParam)->code)
			{
				case LVN_ITEMCHANGED:
				{
					TCHAR szBuffer[MAX_PATH];
					LVITEM lvItem;
					lvItem.mask		= LVIF_TEXT|LVIF_PARAM;
					lvItem.iItem	= ListView_GetSelectionMark(hListView);
					lvItem.iSubItem	= 0;
					lvItem.pszText	= szBuffer;
					lvItem.cchTextMax	= MAX_PATH;
					ListView_GetItem(hListView, &lvItem);

					vector<MODULEENTRY32> vtModuleList;
					RetrieveModules(vtModuleList, (HANDLE)lvItem.lParam);
					
					for (int ii = 0; ii < vtModuleList.size(); ii++)
					{
						TCHAR szTmp[MAX_PATH];
						_stprintf(szTmp, _T("Module Name: %s \tBaseAddr:  0x%06X\t  Size: 0x%06X\t  Module: %s\n"), 
							      vtModuleList[ii].szModule, vtModuleList[ii].modBaseAddr,
								  vtModuleList[ii].modBaseSize, vtModuleList[ii].szExePath);
						OutputDebugString(szTmp);
					}
					
					break;
				}
			}
			break;
		}
		case WM_COMMAND:
		{
			switch (LOWORD(wParam))
			{
				case IDC_BTNEXIT:
				{
					PostQuitMessage(0);
					break;
				}
				case IDC_BTNREFRESH:
				{
					vector<PROCESSENTRY32> vtProcess;
					RetrieveProcess(vtProcess);
					
					ListView_DeleteAllItems(hListView);
					for (int ii = 0; ii < vtProcess.size(); ii++)
					{
						LVITEM lvItem	= {0};
						lvItem.mask		= LVIF_TEXT|LVIF_PARAM;
						lvItem.pszText	= vtProcess[ii].szExeFile;
						lvItem.lParam	= vtProcess[ii].th32ProcessID;
						int nIndex = ListView_InsertItem(hListView, &lvItem);
						TCHAR szTmp[MAX_PATH];
						_stprintf(szTmp, _T("%d"), vtProcess[ii].cntThreads);
						ListView_SetItemText(hListView, nIndex, 1, szTmp);
						_stprintf(szTmp, _T("0X%06X"), vtProcess[ii].th32ProcessID);
						ListView_SetItemText(hListView, nIndex, 2, szTmp);
					}

					break;
				}
			}
			break;
		}
		case WM_CLOSE:
		{
			PostQuitMessage(0);
			break;
		}
	}

	return 0;
}

BOOL WINAPI _tWinMain(HINSTANCE hInstance, HINSTANCE hPreInstance, LPTSTR szCmdLine, int nShowCmd)
{
	DialogBox(hInstance, MAKEINTRESOURCE(IDD_DIALOG), NULL, DlgProc);

	return TRUE;
}